package sample.config;

import brave.Span;
import brave.Tracer;
import brave.Tracing;
import brave.propagation.ThreadLocalSpan;
import brave.propagation.TraceContext;
import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;

/**
 * Created on 2020/7/1
 */
@Intercepts({
		@Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class,
				RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class}),
		@Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class,
				RowBounds.class, ResultHandler.class}),
		@Signature(type = Executor.class, method = "update", args = {MappedStatement.class, Object.class})
})
public class TraceInterceptor implements Interceptor {

	@Override
	public Object intercept(Invocation invocation) throws Throwable {
		Span span = ThreadLocalSpan.CURRENT_TRACER.next();
		if (span == null || span.isNoop()) return invocation.proceed();
		try {
			MappedStatement mappedStatement = (MappedStatement) invocation.getArgs()[0];
			Object parameter = invocation.getArgs()[1];
			// 此处指定CLIENT才可正确显示父子关系
			span.kind(Span.Kind.CLIENT);
			span.name(mappedStatement.getId());
			span.remoteServiceName(mappedStatement.getConfiguration().getEnvironment().getDataSource().getConnection().getCatalog());
			span.tag("sql", mappedStatement.getBoundSql(parameter).getSql());
			span.tag("param", parameter.toString());
			span.start();
			return invocation.proceed();
		} finally {
			Span span1 = ThreadLocalSpan.CURRENT_TRACER.remove();
			span1.finish();
		}
	}
}
